import tensorflow as tf
from tensorflow import keras

layers = tf.keras.layers
Sequential = tf.keras.Sequential
# 优化器
optimizers = tf.keras.optimizers
# 损失函数
losses = tf.keras.losses


class MyDense(layers.Layer):

    def __init__(self, inp_dim, outp_dim):
        super(MyDense, self).__init__()
        # self.kernel = self.add_variable('w', [inp_dim, outp_dim], trainable=False)
        # 通过 tf.Variable 创建的类成员也会自动加入类参数列表
        self.kernel = tf.Variable(tf.random.normal([inp_dim, outp_dim]), trainable=False)

    def call(self, inputs, training=None):
        out = inputs @ self.kernel
        out = tf.nn.relu(out)
        return out


net = MyDense(4, 3)
print(net.variables)
print(net.trainable_variables)
